import asyncio
from datetime import datetime

from common.log import logger
import time
import traceback
from db.es.es_service import ElasticsearchStorage
from bot.insurance_planner_gpt.agent.base_agent import LLMAgent
from db.mysql.entity_service import EntitiesDao
import settings

entities_dao = EntitiesDao()

default_template = """
你是一个AI助手，你的任务是以用户的视角把对话中用户的需求归纳成一句话，字数限制在15个以内。突出用户的需求让用户可以清晰的分辨本次对话的核心：
示例：
助手：嗨！您好！👋我是犀心小助，您的专属保险规划师。您在保险规划方面有什么需要我帮忙的吗？我将竭诚为您提供个性化的保险规划服务哦！😊
用户：我想了解下个人保险规划 主要是医疗保障方面的
助手：您好！非常明白您的需求，您想要了解个人的医疗保险规划。😊 这是一个很重要的问题，毕竟，良好的医疗保险规划，能给您带来更加全面的医疗保障。
为了更好地帮助您进行个人的医疗保险规划，我需要了解一下您的一些具体情况，例如👇：
1. 您希望医疗保险的预算大概是多少？
2. 您对医疗保险有哪些具体的需求或期待？
3. 您可能会面临的具体医疗风险有哪些？
这样，我才能为您规划出更适合您的个人医疗保险方案哦！期待您的回复！😉
输出：用户需要全面的医疗保障咨询
示例结束
===
{document}
===
请输出你总结的内容：
        """


class ConversationUtil:
    def __init__(self, ):
        self.es = ElasticsearchStorage()

    def store_conversation_index_(self, index_name, session_id, message_id, msg_type, text_info, role):
        try:
            conversation = {
                "session_id": session_id,
                "role": role,
                "message_id": message_id,
                "attitude": "no",
                "msg_type": msg_type,
                "text_info": text_info,
                "message_time": datetime.now()
            }
            self.es.index_document(index_name, conversation)
        except Exception as ee:
            logger.error(traceback.format_exc())

    def store_qw_conversation(self, session_id, message_id, msg_type, text_info, role):
        self.store_conversation_index_(settings.QW_CONVERSATION_HISTORY_INDEX, session_id, message_id, msg_type,
                                       text_info, role)

    def store_qw_conversation_new(self, conversations):

        for conversation in conversations:
            self.es.index_document(settings.QW_CONVERSATION_INDEX, conversation)

    def store_add_wechat_conversation(self, conversations):
        # TODO 改成批量插入
        # doc = []
        # for conversation in conversations:
        #     doc.append({"_index": 'add_wechat', '_op_type': 'index', "_type": "doc", "id": "", "_source": conversation})
        # success, _ = self.es.bulk(doc)

        for conversation in conversations:
            self.es.index_document('add_wechat', conversation)

    def store_conversation_index(self, session_id, message_id, msg_type, text_info, role,
                                 client_type="planner", receive_time=None):
        try:
            if session_id and session_id[:4] == "LXKJ":
                return
            conversation = {
                "session_id": session_id,
                "role": role,
                "message_id": message_id,
                "attitude": "no",
                "msg_type": msg_type,
                "text_info": text_info,
                "message_time": receive_time if receive_time else datetime.now(),
                "client_type": client_type
            }
            self.es.index_document(settings.CONVERSATION_HISTORY_INDEX, conversation)
            entities_dao.insert_ql_qw_contact_msg(conversation)
        except Exception as ee:
            logger.error(traceback.format_exc())

    def store_planner_qw_conversation_index(self, session_id, message_id, msg_type, text_info, role):
        """ 存储规划师企微对话数据 """
        try:
            if session_id and session_id[:4] == "LXKJ":
                return
            conversation = {
                "session_id": session_id,
                "role": role,
                "message_id": message_id,
                "attitude": "no",
                "msg_type": msg_type,
                "text_info": text_info,
                "message_time": datetime.now(),
                "client_type": "planner_qw"
            }
            self.es.index_document(settings.PLANNER_QW_CONVERSATION_HISTORY_INDEX, conversation)
        except Exception as ee:
            logger.error(traceback.format_exc())

    def store_conversation_user(self, session_id, user_name, conversationHistory):
        # 用户<->会话关系存储
        session_title = None
        logger.info('save to es——{},{}'.format(user_name, session_id))
        if not conversationHistory:
            session_title = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        else:
            self.summarize_bot = ConversationToSummarize(f'''{conversationHistory}''')
            session_title = asyncio.run(self.summarize_bot.achat_auto_llm(type="zephyr"))
            logger.info('user：{},session_title:{}'.format(user_name, session_title))

        try:
            query = {
                "query": {
                    "bool": {
                        "must": [
                            {"term":
                                 {"session_id": session_id}
                             },
                            {"term":
                                 {"user_name": user_name}
                             }
                        ]
                    }
                }
            }
            results = self.es.search_documents(settings.USER_SESSION_INDEX, query)
            if not results:
                logger.info(
                    'create new user-sessionId-session_title：{}-{}'.format(user_name, session_id, session_title))
                # 创建用户会话
                conversation = {
                    "user_name": user_name,
                    "session_id": session_id,
                    "session_title": session_title,
                    "message_time": datetime.now()
                }
                self.es.index_document(settings.USER_SESSION_INDEX, conversation)
            elif conversationHistory and session_title:
                logger.info('update title')
                # 更新会话标题
                updated_fields = {
                    'session_title': session_title,
                    "message_time": datetime.now()
                }
                self.es.update_document(settings.USER_SESSION_INDEX, results[0]['_id'], updated_fields)
        except Exception as ee:
            logger.error(traceback.format_exc())

    def get_dialogue_history_by_session_id(self, session_id):
        query = {
            "query": {
                "bool": {
                    "must": [
                        {
                            "term": {
                                "session_id": {
                                    "value": session_id
                                }
                            }
                        }
                    ]
                }
            }, "sort": [
                {
                    "message_time": {
                        "order": "asc"
                    }
                }
            ]
            , "size": 10000
        }
        results = self.es.search_documents(settings.CONVERSATION_HISTORY_INDEX, query)
        messages = []
        if results:
            delete = []
            for hit in results:
                try:
                    message_id = hit["_source"]["message_id"]
                    # 展示历史记录的时候只保留坐席修改后的对话
                    if 'agent_' in message_id:
                        delete.append('origin_' + message_id.replace('agent_', ''))
                    msg_type = hit["_source"].get('msg_type', 'text')

                    role = hit["_source"]["role"]
                    timestamp = hit["_source"]["message_time"]
                    text_info = hit["_source"]["text_info"]
                    # timestamp = datetime.strftime(datetime.strptime(timestamp, '%Y-%m-%dT%H:%M:%S.%f'),
                    #                               '%Y-%m-%d %H:%M:%S')
                    content = {}
                    content['role'] = role
                    content['msg_type'] = msg_type
                    content['msg_time'] = timestamp
                    content['msg_id'] = message_id
                    content['content'] = text_info
                    messages.append(content)

                except Exception:
                    logger.error(traceback.format_exc())
                    continue

            messages = [message for message in messages if
                        message.get('message_id') not in delete]
        return messages


class ConversationToSummarize(LLMAgent):
    def __init__(self, document):
        self.prompt = default_template
        self.prompt = self.prompt.format(document=document)
        super().__init__(self.prompt)


if __name__ == '__main__':
    cu = ConversationUtil()
    print(cu.get_dialogue_history_by_session_id("userFZaJYymJxBUMpifgxjIJ:userCbZn5NZsGMlsInm7OOvn"))
